# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import math
import torch.nn.functional as F
import torch
from torch import autograd

from fairseq import utils

from fairseq.criterions import FairseqCriterion, register_criterion


@register_criterion('story_coherence_enhance_fix')
class StoryCoherenceEF(FairseqCriterion):

    def __init__(self, args, task):
        super().__init__(args, task)

    def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        with autograd.detect_anomaly():
            net_output = model(**sample['net_input'])
            neg_1 = net_output['neg_1']
            neg_2 = net_output['neg_2']
            sentence_pair_acc = net_output['sentence_pair_acc']

            loss, nll_loss, article_pair_acc, pos_scores, neg_scores = self.compute_loss(model, net_output, sample, reduce=reduce)
            sample_size = 1
            logging_output = {
                'loss': utils.item(loss.data) if reduce else loss.data,
                'nll_loss': utils.item(nll_loss.data) if reduce else nll_loss.data,
                'ntokens': sample['ntokens'],
                'nsentences': sample['target'].size(0),
                'sample_size': sample_size,
                'article_pair_acc': article_pair_acc.item(),
                'pos_scores': pos_scores.item(),
                'neg_scores': neg_scores.item(),
                'neg_1': neg_1.item(),
                'neg_2': neg_2.item(),
                'sentence_pair_acc': sentence_pair_acc.item()

            }
            return loss, sample_size, logging_output

    def compute_loss(self, model, net_output, sample, reduce=True):
        pos_scores = net_output['pos_scores']
        neg_scores = net_output['neg_scores']
        score_pad_mask = net_output['score_pad_mask']


        # mean of positive sentence pair, mean of negative sentence pair
        pos_scores = pos_scores[score_pad_mask]
        neg_scores = neg_scores[score_pad_mask]
        #
        # margin loss
        scores = 0.5 - pos_scores + neg_scores
        scores = scores.clamp(min=0)
        nll_loss = loss = scores.mean()

        # # bce loss
        # pos_scores = pos_scores[score_pad_mask]
        # neg_scores = neg_scores[score_pad_mask]
        # loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores), reduction='mean')
        # loss += F.binary_cross_entropy(neg_scores, torch.zeros_like(neg_scores), reduction='mean')
        # nll_loss = loss

        # compute article accuracy
        pos_article_scores = net_output['pos_article_scores']
        neg_article_scores = net_output['neg_article_scores']

        article_pair_acc = (pos_article_scores > neg_article_scores).float().mean()

        return loss, nll_loss, article_pair_acc, pos_scores.mean(), neg_scores.mean()

    @staticmethod
    def aggregate_logging_outputs(logging_outputs):
        """Aggregate logging outputs from data parallel training."""
        ntokens = sum(log.get('ntokens', 0) for log in logging_outputs)
        nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)
        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)
        return {
            'loss': sum(log.get('loss', 0) for log in logging_outputs) / sample_size / math.log(2),
            'nll_loss': sum(log.get('nll_loss', 0) for log in logging_outputs) / sample_size / math.log(2),
            'ntokens': ntokens,
            'nsentences': nsentences,
            'sample_size': sample_size,
            'article_pair_acc' : sum(log.get('article_pair_acc', 0) for log in logging_outputs) / len(logging_outputs),
            'pos_scores': sum(log.get('pos_scores', 0) for log in logging_outputs) / len(logging_outputs),
            'neg_scores': sum(log.get('neg_scores', 0) for log in logging_outputs) / len(logging_outputs),
            'neg_1': sum(log.get('neg_1', 0) for log in logging_outputs) / len(logging_outputs),
            'neg_2': sum(log.get('neg_2', 0) for log in logging_outputs) / len(logging_outputs),
            'sentence_pair_acc': sum(log.get('sentence_pair_acc', 0) for log in logging_outputs) / len(logging_outputs),
        }
